from __future__ import print_function, division
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
import torch
from SwinTransformer.models.swin_transformer import ViTSwinTransformerBlock, VitPatchMerging
from einops import rearrange

class conv_block(nn.Module):
    """
    Convolution Block
    """

    def __init__(self, in_ch, out_ch):
        super(conv_block, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),

            nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True))

    def forward(self, x):
        x = self.conv(x)

        return x


class up_conv(nn.Module):
    """
    Up Convolution Block
    """

    def __init__(self, in_ch, out_ch, scale_factor=2):
        super(up_conv, self).__init__()
        self.up = nn.Sequential(
            nn.Upsample(scale_factor=scale_factor),
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x


class ConvBNRelu1D(nn.Module):
    def __init__(self, inChannel, outChannel, kernel_size=3, padding=1, bias=False):
        super(ConvBNRelu1D, self).__init__()
        self.conv = nn.Conv1d(inChannel, outChannel, kernel_size=kernel_size, padding=padding, bias=bias)
        self.bn = nn.BatchNorm1d(outChannel, momentum=0.1, affine=True)
        self.reLu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return self.reLu(x)




class mutiDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose1_0 = up_conv(depths[1], depths[0])
        self.Conv1_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose2_1 = up_conv(depths[2], depths[1])
        self.Conv2_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose2_0 = up_conv(depths[1], depths[0])
        self.Conv2_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose3_2 = up_conv(depths[3], depths[2])
        self.Conv3_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose3_1 = up_conv(depths[2], depths[1])
        self.Conv3_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose3_0 = up_conv(depths[1], depths[0])
        self.Conv3_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3])
        self.Conv4_3 = conv_block(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv(depths[3], depths[2])
        self.Conv4_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv(depths[2], depths[1])
        self.Conv4_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv(depths[1], depths[0])
        self.Conv4_0 = conv_block(depths[0], depths[0])

        self.ConvAll = conv_block(depths[0] * 5, depths[0])
        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        level0 = levels[0]
        # level1的解码
        level1 = self.ConvTranspose1_0(levels[1])
        level1 = self.Conv1_0(level1)

        # level2的解码
        level2 = self.ConvTranspose2_1(levels[2])
        level2 = torch.cat([level2, levels[1]], dim=1)
        level2 = self.Conv2_1(level2)

        level2 = self.ConvTranspose2_0(level2)
        level2 = self.Conv2_0(level2)

        # level3的解码
        level3 = self.ConvTranspose3_2(levels[3])
        level3 = torch.cat([level3, levels[2]], dim=1)
        level3 = self.Conv3_2(level3)

        level3 = self.ConvTranspose3_1(level3)
        level3 = torch.cat([level3, levels[1]], dim=1)
        level3 = self.Conv3_1(level3)

        level3 = self.ConvTranspose3_0(level3)
        level3 = self.Conv3_0(level3)

        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        # 所有level的融合
        out = torch.cat([level0, level1, level2, level3, level4], dim=1)
        out = self.ConvAll(out)
        out = self.Out(out)

        return out


class UnetDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3])
        self.Conv4_3 = conv_block(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv(depths[3], depths[2])
        self.Conv4_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv(depths[2], depths[1])
        self.Conv4_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv(depths[1], depths[0])
        self.Conv4_0 = conv_block(depths[0], depths[0])

        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        out = self.Out(level4)

        return out


class DenseDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose1_0 = up_conv(depths[1], depths[0], scale_factor=2)

        self.ConvTranspose2_1 = up_conv(depths[2], depths[1], scale_factor=2)

        self.ConvTranspose2_0 = up_conv(depths[2], depths[0], scale_factor=4)

        self.ConvTranspose3_2 = up_conv(depths[3], depths[2], scale_factor=2)

        self.ConvTranspose3_1 = up_conv(depths[3], depths[1], scale_factor=4)

        self.ConvTranspose3_0 = up_conv(depths[3], depths[0], scale_factor=8)

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3], scale_factor=2)

        self.ConvTranspose4_2 = up_conv(depths[4], depths[2], scale_factor=4)

        self.ConvTranspose4_1 = up_conv(depths[4], depths[1], scale_factor=8)

        self.ConvTranspose4_0 = up_conv(depths[4], depths[0], scale_factor=16)

        self.upConv3 = conv_block(depths[3] * 2, depths[3])
        self.upConv2 = conv_block(depths[2] * 3, depths[2])
        self.upConv1 = conv_block(depths[1] * 4, depths[1])
        self.upConv0 = conv_block(depths[0] * 5, depths[0])

        self.out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        l4_3 = self.ConvTranspose4_3(levels[4])
        l4_2 = self.ConvTranspose4_2(levels[4])
        l4_1 = self.ConvTranspose4_1(levels[4])
        l4_0 = self.ConvTranspose4_0(levels[4])

        levels[3] = torch.cat([levels[3], l4_3], dim=1)
        levels[3] = self.upConv3(levels[3])

        l3_2 = self.ConvTranspose3_2(levels[3])
        l3_1 = self.ConvTranspose3_1(levels[3])
        l3_0 = self.ConvTranspose3_0(levels[3])

        levels[2] = torch.cat([levels[2], l4_2, l3_2], dim=1)
        levels[2] = self.upConv2(levels[2])

        l2_1 = self.ConvTranspose2_1(levels[2])
        l2_0 = self.ConvTranspose2_0(levels[2])

        levels[1] = torch.cat([levels[1], l4_1, l3_1, l2_1], dim=1)
        levels[1] = self.upConv1(levels[1])

        l1_0 = self.ConvTranspose1_0(levels[1])

        levels[0] = torch.cat([levels[0], l4_0, l3_0, l2_0, l1_0], dim=1)
        levels[0] = self.upConv0(levels[0])

        out = self.out(levels[0])

        return out









class depthwise_separable_conv(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1, kernel_size=3, padding=1, bias=False):
        super().__init__()
        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size=kernel_size, padding=padding, groups=in_ch, bias=bias, stride=stride)
        self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=bias)

    def forward(self, x):
        out = self.depthwise(x)
        out = self.pointwise(out)

        return out
class RelativePositionBias(nn.Module):
    # input-independent relative position attention
    # As the number of parameters is smaller, so use 2D here
    # Borrowed some code from SwinTransformer: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
    def __init__(self, num_heads, h, w):
        super().__init__()
        self.num_heads = num_heads
        self.h = h
        self.w = w

        self.relative_position_bias_table = nn.Parameter(
            torch.randn((2 * h - 1) * (2 * w - 1), num_heads) * 0.02)

        coords_h = torch.arange(self.h)
        coords_w = torch.arange(self.w)
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, h, w
        coords_flatten = torch.flatten(coords, 1)  # 2, hw

        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += self.h - 1
        relative_coords[:, :, 1] += self.w - 1
        relative_coords[:, :, 0] *= 2 * self.h - 1
        relative_position_index = relative_coords.sum(-1)  # hw, hw

        self.register_buffer("relative_position_index", relative_position_index)

    def forward(self, H, W):
        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(self.h,
                                                                                                               self.w,
                                                                                                               self.h * self.w,
                                                                                                               -1)  # h, w, hw, nH
        relative_position_bias_expand_h = torch.repeat_interleave(relative_position_bias, H // self.h, dim=0)
        relative_position_bias_expanded = torch.repeat_interleave(relative_position_bias_expand_h, W // self.w,
                                                                  dim=1)  # HW, hw, nH

        relative_position_bias_expanded = relative_position_bias_expanded.view(H * W, self.h * self.w,
                                                                               self.num_heads).permute(2, 0,
                                                                                                       1).contiguous().unsqueeze(
            0)

        return relative_position_bias_expanded

class LinearAttention(nn.Module):

    def __init__(self, dim, heads=4, dim_head=8, attn_drop=0., proj_drop=0., reduce_size=16, projection='interp',
                 rel_pos=True, H=256):
        super().__init__()

        self.inner_dim = dim_head * heads
        self.heads = heads
        self.scale = dim_head ** (-0.5)
        self.dim_head = dim_head
        self.reduce_size = reduce_size
        self.projection = projection
        self.rel_pos = rel_pos

        # depthwise conv is slightly better than conv1x1
        # self.to_qkv = nn.Conv2d(dim, self.inner_dim*3, kernel_size=1, stride=1, padding=0, bias=True)
        # self.to_out = nn.Conv2d(self.inner_dim, dim, kernel_size=1, stride=1, padding=0, bias=True)

        self.to_qkv = depthwise_separable_conv(dim, self.inner_dim * 3)
        self.to_out = depthwise_separable_conv(self.inner_dim, dim)

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj_drop = nn.Dropout(proj_drop)

        if self.rel_pos:
            # 2D input-independent relative position encoding is a little bit better than
            # 1D input-denpendent counterpart
            self.relative_position_encoding = RelativePositionBias(heads, reduce_size, reduce_size)
            # self.relative_position_encoding = RelativePositionEmbedding(dim_head, reduce_size)
        self.reduceK = nn.Linear(H * H, reduce_size * reduce_size)
        self.reduceV = nn.Linear(H * H, reduce_size * reduce_size)

    def forward(self, x):

        B, C, H, W = x.shape

        # B, inner_dim, H, W
        qkv = self.to_qkv(x)
        q, k, v = qkv.chunk(3, dim=1)

        if self.projection == 'interp' and H != self.reduce_size:
            k, v = map(lambda t: F.interpolate(t, size=self.reduce_size, mode='bilinear', align_corners=True), (k, v))

        elif self.projection == 'maxpool' and H != self.reduce_size:
            k, v = map(lambda t: F.adaptive_max_pool2d(t, output_size=self.reduce_size), (k, v))

        elif self.projection == 'learn' and H != self.reduce_size:
            k=k.reshape(B,self.inner_dim,H*W)
            k = self.reduceK(k)
            k = k.reshape(B, self.inner_dim, self.reduce_size, self.reduce_size)

            v = v.reshape(B, self.inner_dim, H * W)
            v = self.reduceV(v)
            v = v.reshape(B, self.inner_dim, self.reduce_size, self.reduce_size)

        q = rearrange(q, 'b (dim_head heads) h w -> b heads (h w) dim_head', dim_head=self.dim_head, heads=self.heads,
                      h=H, w=W)
        k, v = map(lambda t: rearrange(t, 'b (dim_head heads) h w -> b heads (h w) dim_head', dim_head=self.dim_head,
                                       heads=self.heads, h=self.reduce_size, w=self.reduce_size), (k, v))

        q_k_attn = torch.einsum('bhid,bhjd->bhij', q, k)

        if self.rel_pos:
            relative_position_bias = self.relative_position_encoding(H, W)
            q_k_attn += relative_position_bias
            # rel_attn_h, rel_attn_w = self.relative_position_encoding(q, self.heads, H, W, self.dim_head)
            # q_k_attn = q_k_attn + rel_attn_h + rel_attn_w

        q_k_attn *= self.scale
        q_k_attn = F.softmax(q_k_attn, dim=-1)
        q_k_attn = self.attn_drop(q_k_attn)

        out = torch.einsum('bhij,bhjd->bhid', q_k_attn, v)
        out = rearrange(out, 'b heads (h w) dim_head -> b (dim_head heads) h w', h=H, w=W, dim_head=self.dim_head,
                        heads=self.heads)

        out = self.to_out(out)
        out = self.proj_drop(out)

        return out


class edgeNet(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=2):
        super().__init__()

        n1 = 32
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.local=nn.Sequential(
            conv_block(3, filters[0]),
            conv_block(filters[0], filters[0]),
            conv_block(filters[0], filters[0]),
            conv_block(filters[0], filters[0]),
            conv_block(filters[0], filters[0]),

                                 )


        self.Conv1 = conv_block(in_ch, filters[0])
        # self.Local=nn.Sequential(
        #     conv_block(in_ch, filters[0]),
        #     conv_block(filters[0], filters[1]),
        #     conv_block(filters[1], filters[2]),
        #     conv_block(filters[2], filters[3]),
        #     conv_block(filters[3], filters[4])
        #
        # )
        # self.active = torch.nn.Sigmoid()
        self.gl= nn.Sequential(



            LinearAttention(dim=filters[0], H=64, projection="interp", reduce_size=8),
            LinearAttention(dim=filters[0], H=64, projection="interp", reduce_size=8),
            LinearAttention(dim=filters[0], H=64, projection="interp", reduce_size=8),
            LinearAttention(dim=filters[0], H=64, projection="interp", reduce_size=8),
                                       )


        # self.UNetDecoder = UnetDecoder(filters)
        self.out=nn.Conv2d(filters[0], out_ch, kernel_size=1, stride=1, padding=0)
    def forward(self, input):
        g1 = self.Conv1(input)
        g1=self.gl(g1)

        e1=self.local(input)



        out = self.out(g1+e1)

        return out